#!/usr/bin/python
# -*- coding: UTF-8 -*-

from PIL import Image, ImageDraw
import os
import random
import math
import multiprocessing
import xml.dom.minidom

def augmentation_resize(inputfolder, outputfolder, tmpfolder, filename, resizeAugmentTimes):
    '''
    对指定目录下的指定名字的jpg和xml文件进行缩放形式的数据增广，并保存到指定目录下
    :param inputfolder:
    :param outputfolder:
    :param filename:
    :return:
    '''
    #打开图片
    inputfilepath_jpg = inputfolder + '/' + filename + '.jpg'
    im = Image.open(inputfilepath_jpg)

    #打开标注数据xml文件
    inputfilepath_xml = inputfolder + '/' + filename + '.xml'
    DOMTree = xml.dom.minidom.parse(inputfilepath_xml)
    collection = DOMTree.documentElement
    width = collection.getElementsByTagName("width")
    height = collection.getElementsByTagName("height")
    xmlxmin = collection.getElementsByTagName("xmin")
    xmlxmax = collection.getElementsByTagName("xmax")
    xmlymin = collection.getElementsByTagName("ymin")
    xmlymax = collection.getElementsByTagName("ymax")
    rectNum = len(xmlxmin)
    w = float(width[0].childNodes[0].data)
    h = float(height[0].childNodes[0].data)

    #错误检查
    #if(len(xmlfolder) == 0 or len(xmlfilename) == 0 or len(xmlpath) == 0 or len(xmlxmin) == 0):
    #    return

    #图片和数据进行增广：方法采用等比例缩放
    index = 0
    for k in range(int(10), int(10+resizeAugmentTimes), 1):
        resizecoef = 0.05 + k / 10
        outputfilepath_jpg = outputfolder + '/' + filename + '_' + str(index) + '.jpg'
        tmpfilepath_jpg = tmpfolder + '/' + filename + '_' + str(index) + '.jpg'
        outputfilepath_xml = outputfolder + '/' + filename + '_' + str(index) + '.xml'
        DOMTreeClone = DOMTree.cloneNode(10)

        # 检查boundary box越界
        ifcrossflag = 0
        halfdwid = int((w * resizecoef - w) / 2.0)
        halfdhei = int((h * resizecoef - h) / 2.0)
        for i in range(rectNum):
            xmin = float(xmlxmin[i].childNodes[0].data) * resizecoef
            xmax = float(xmlxmax[i].childNodes[0].data) * resizecoef
            ymin = float(xmlymin[i].childNodes[0].data) * resizecoef
            ymax = float(xmlymax[i].childNodes[0].data) * resizecoef
            if xmin <= halfdwid + 1:
                ifcrossflag = 1
                break
            if xmax >= (w + halfdwid - 1):
                ifcrossflag = 1
                break
            if ymin <= halfdhei + 1:
                ifcrossflag = 1
                break
            if ymax >= (h + halfdhei - 1):
                ifcrossflag = 1
                break
        if(ifcrossflag == 1):  # 图像放大后裁切，若标注框超出边界，则停止继续放大
            resizecoef = 1.0

        if resizecoef > 1.0: #图像放大，则需要裁切
            ##图像放大后裁切并保存
            halfdwid = int((w * resizecoef - w) / 2.0)
            halfdhei = int((h * resizecoef - h) / 2.0)
            im_resized = im.resize((int(w * resizecoef), int(h * resizecoef)))
            region = im_resized.crop((halfdwid, halfdhei, int(w + halfdwid), int(h + halfdhei)))
            region.save(outputfilepath_jpg, "JPEG")

            # 图像的标注数据同步变动保存
            DOMTreeClone.documentElement.getElementsByTagName("filename")[0].childNodes[0].data = filename + '_' + str(index) + '.jpg'
            #DOMTreeClone.documentElement.getElementsByTagName("path")[0].childNodes[0].data = outputfilepath_jpg
            DOMTreeClone.documentElement.getElementsByTagName("width")[0].childNodes[0].data = str(int(w))
            DOMTreeClone.documentElement.getElementsByTagName("height")[0].childNodes[0].data = str(int(h))
            for i in range(rectNum):
                DOMTreeClone.documentElement.getElementsByTagName("xmin")[i].childNodes[0].data = str(int(float(xmlxmin[i].childNodes[0].data) * resizecoef - halfdwid))
                DOMTreeClone.documentElement.getElementsByTagName("xmax")[i].childNodes[0].data = str(int(float(xmlxmax[i].childNodes[0].data) * resizecoef - halfdwid))
                DOMTreeClone.documentElement.getElementsByTagName("ymin")[i].childNodes[0].data = str(int(float(xmlymin[i].childNodes[0].data) * resizecoef - halfdhei))
                DOMTreeClone.documentElement.getElementsByTagName("ymax")[i].childNodes[0].data = str(int(float(xmlymax[i].childNodes[0].data) * resizecoef - halfdhei))
            with open(outputfilepath_xml, 'w', encoding='utf-8') as f:
                DOMTreeClone.writexml(f, encoding='utf-8')

            #画框用于测试效果
            '''draw = ImageDraw.Draw(region)
            for i in range(rectNum):
                xl = float(DOMTreeClone.documentElement.getElementsByTagName("xmin")[i].childNodes[0].data)
                xr = float(DOMTreeClone.documentElement.getElementsByTagName("xmax")[i].childNodes[0].data)
                yu = float(DOMTreeClone.documentElement.getElementsByTagName("ymin")[i].childNodes[0].data)
                yd = float(DOMTreeClone.documentElement.getElementsByTagName("ymax")[i].childNodes[0].data)
                draw.line([(xl, yu), (xr, yu)], "red")
                draw.line([(xr, yu), (xr, yd)], "red")
                draw.line([(xr, yd), (xl, yd)], "red")
                draw.line([(xl, yd), (xl, yu)], "red")
            region.save(tmpfilepath_jpg, "JPEG")'''
        index = index + 1
    return

def findfiles(path, t):
    '''
    遍历查找目录下的所有jpg文件名
    :param path:
    :param t:
    :return:
    '''
    files = os.listdir(path)
    for f in files:
        npath = path + '/' + f
        if(os.path.isfile(npath)):
            if(os.path.splitext(npath)[1] ==".jpg"):
                t.append(os.path.splitext(f)[0])
                print(os.path.splitext(f)[0])
    return

if __name__=="__main__":
    '''
    主函数
    输入：带标注信息的图像
    输出：图像缩放后的一组图像，以及同步得到的对应标注信息，放入outputfiles下          
    note: 数据增广的总倍数 = resizeAugmentTimes
    '''
    print("version:2018.11.24 auther:lhf")
    resizeAugmentTimes = 5  # 通过缩放方式进行数据增广的倍数

    # 图像缩放后的一组图像，以及同步得到的对应标注信息，放入outputfiles下
    inputfolder = "./trainData"
    outputfolder = "./outputFiles"
    tmpfolder = "./tmpFiles"

    filenamelist = []
    inputfolder_copy = inputfolder
    findfiles(inputfolder_copy, filenamelist)
    print("findfiles finish")

    # 按照缩放的方式对原始图像及标注数据进行增广
    pool = multiprocessing.Pool(processes=multiprocessing.cpu_count())
    for filename in filenamelist:
        outputfolder_copy = outputfolder
        tmpfolder_copy = tmpfolder
        if (os.path.exists(outputfolder_copy)):
            #print("111")
            x = 1
        else:
            os.makedirs(outputfolder_copy)
            #os.makedirs(tmpfolder_copy)
            print("mkdir" + outputfolder_copy)
            #print("mkdir" + tmpfolder_copy)

        #print("processing %s ..." % filename)
        pool.apply_async(augmentation_resize, (inputfolder_copy, outputfolder_copy, tmpfolder_copy, filename, resizeAugmentTimes))
    pool.close()
    pool.join()
    filenamelist.clear()
